import torch
import numpy as np
import cv2
from efficient_track_anything.automatic_mask_generator import (
    EfficientTAMAutomaticMaskGenerator,
)
from efficient_track_anything.build_efficienttam import build_efficienttam


class Efficient_SAM_Encoder():

    def __init__(self, sam_ckpt_path,model_cfg):
        self.device = 'cuda:0'
        self.sam = build_efficienttam(model_cfg, sam_ckpt_path, device=self.device)
        # 参数检查是否和langspalt一致
        self.mask_generator = EfficientTAMAutomaticMaskGenerator(
            model=self.sam,
            points_per_side=32,
            pred_iou_thresh=0.7,
            box_nms_thresh=0.7,
            stability_score_thresh=0.85,
            crop_n_layers=1,
            crop_n_points_downscale_factor=1,
            min_mask_region_area=100,
        )

    def mask2segmap(self,masks, image):
            seg_img_list = []
            #seg_map 全为-1，shape=H,W
            seg_map = -np.ones(image.shape[:2], dtype=np.int32)
            for i in range(len(masks)):
                mask = masks[i]
                #根据mask从图像中提取出掩膜部分的图像
                seg_img = self.get_seg_img(mask, image)
                #不失resize成224*224*3
                pad_seg_img = cv2.resize(self.pad_img(seg_img), (224,224))
                seg_img_list.append(pad_seg_img)

                seg_map[masks[i]['segmentation']] = i
            seg_imgs = np.stack(seg_img_list, axis=0) # b,H,W,3
            seg_imgs = (torch.from_numpy(seg_imgs.astype("float32")).permute(0,3,1,2) / 255.0).to('cuda')
            #seg_imgs b,3,H,W     seg_maps H,W
            return seg_imgs, seg_map

    def sam_encoder(self,image):
        #image 1,C,H,W -> H,W,C 
        image = cv2.cvtColor(image[0].permute(1,2,0).numpy().astype(np.uint8), cv2.COLOR_BGR2RGB)
        # pre-compute masks
        #masks_default, masks_s, masks_m, masks_l = self.mask_generator.generate(image)
        masks = self.mask_generator.generate(image)
        # pre-compute postprocess
        masks = self.masks_update(masks, iou_thr=0.8, score_thr=0.7, inner_thr=0.5)[0]
        #seg_images, seg_maps = {}, {}
        #seg_imgs b,3,H,W     seg_maps H,W
        seg_images, seg_maps = self.mask2segmap(masks, image)
        '''
        if len(masks_s) != 0:
            seg_images['s'], seg_maps['s'] = self.mask2segmap(masks_s, image)
        if len(masks_m) != 0:
            seg_images['m'], seg_maps['m'] = self.mask2segmap(masks_m, image)
        if len(masks_l) != 0:
            seg_images['l'], seg_maps['l'] = self.mask2segmap(masks_l, image)
        '''
        return seg_images, seg_maps

    def get_seg_img(self,mask, image):
        #image (H, W, C)
        image = image.copy()
        #这个时候image在mask为0的地方会被值为0的
        image[mask['segmentation']==0] = np.array([0, 0,  0], dtype=np.uint8)
        x,y,w,h = np.int32(mask['bbox'])
        #切出分割出来的部分
        seg_img = image[y:y+h, x:x+w, ...]
        return seg_img

    def pad_img(self,img):
        h, w, _ = img.shape
        l = max(w,h)
        pad = np.zeros((l,l,3), dtype=np.uint8)
        if h > w:
            pad[:,(h-w)//2:(h-w)//2 + w, :] = img
        else:
            pad[(w-h)//2:(w-h)//2 + h, :, :] = img
        return pad

    def masks_update(self,*args, **kwargs):
        # remove redundant masks based on the scores and overlap rate between masks
        masks_new = ()
        for masks_lvl in (args):
            seg_pred =  torch.from_numpy(np.stack([m['segmentation'] for m in masks_lvl], axis=0))
            iou_pred = torch.from_numpy(np.stack([m['predicted_iou'] for m in masks_lvl], axis=0))
            stability = torch.from_numpy(np.stack([m['stability_score'] for m in masks_lvl], axis=0))

            scores = stability * iou_pred
            keep_mask_nms = self.mask_nms(seg_pred, scores, **kwargs)
            masks_lvl = self.filter(keep_mask_nms, masks_lvl)

            masks_new += (masks_lvl,)
        return masks_new

    def filter(self,keep: torch.Tensor, masks_result) -> None:
        keep = keep.int().cpu().numpy()
        result_keep = []
        for i, m in enumerate(masks_result):
            if i in keep: result_keep.append(m)
        return result_keep

    def mask_nms(self,masks, scores, iou_thr=0.7, score_thr=0.1, inner_thr=0.2, **kwargs):
        """
    Perform mask non-maximum suppression (NMS) on a set of masks based on their scores.
    
    Args:
        masks (torch.Tensor): has shape (num_masks, H, W)
        scores (torch.Tensor): The scores of the masks, has shape (num_masks,)
        iou_thr (float, optional): The threshold for IoU.
        score_thr (float, optional): The threshold for the mask scores.
        inner_thr (float, optional): The threshold for the overlap rate.
        **kwargs: Additional keyword arguments.
    Returns:
        selected_idx (torch.Tensor): A tensor representing the selected indices of the masks after NMS.
    """

        scores, idx = scores.sort(0, descending=True)
        num_masks = idx.shape[0]
    
        masks_ord = masks[idx.view(-1), :]
        masks_area = torch.sum(masks_ord, dim=(1, 2), dtype=torch.float)

        iou_matrix = torch.zeros((num_masks,) * 2, dtype=torch.float, device=masks.device)
        inner_iou_matrix = torch.zeros((num_masks,) * 2, dtype=torch.float, device=masks.device)
        for i in range(num_masks):
            for j in range(i, num_masks):
                intersection = torch.sum(torch.logical_and(masks_ord[i], masks_ord[j]), dtype=torch.float)
                union = torch.sum(torch.logical_or(masks_ord[i], masks_ord[j]), dtype=torch.float)
                iou = intersection / union
                iou_matrix[i, j] = iou
                # select mask pairs that may have a severe internal relationship
                if intersection / masks_area[i] < 0.5 and intersection / masks_area[j] >= 0.85:
                    inner_iou = 1 - (intersection / masks_area[j]) * (intersection / masks_area[i])
                    inner_iou_matrix[i, j] = inner_iou
                if intersection / masks_area[i] >= 0.85 and intersection / masks_area[j] < 0.5:
                    inner_iou = 1 - (intersection / masks_area[j]) * (intersection / masks_area[i])
                    inner_iou_matrix[j, i] = inner_iou

        iou_matrix.triu_(diagonal=1)
        iou_max, _ = iou_matrix.max(dim=0)
        inner_iou_matrix_u = torch.triu(inner_iou_matrix, diagonal=1)
        inner_iou_max_u, _ = inner_iou_matrix_u.max(dim=0)
        inner_iou_matrix_l = torch.tril(inner_iou_matrix, diagonal=1)
        inner_iou_max_l, _ = inner_iou_matrix_l.max(dim=0)
    
        keep = iou_max <= iou_thr
        keep_conf = scores > score_thr
        keep_inner_u = inner_iou_max_u <= 1 - inner_thr
        keep_inner_l = inner_iou_max_l <= 1 - inner_thr
    
        # If there are no masks with scores above threshold, the top 3 masks are selected
        if keep_conf.sum() == 0:
            index = scores.topk(3).indices
            keep_conf[index, 0] = True
        if keep_inner_u.sum() == 0:
            index = scores.topk(3).indices
            keep_inner_u[index, 0] = True
        if keep_inner_l.sum() == 0:
            index = scores.topk(3).indices
            keep_inner_l[index, 0] = True
        keep *= keep_conf
        keep *= keep_inner_u
        keep *= keep_inner_l

        selected_idx = idx[keep]
        return selected_idx